#include "hla_processor.h"

static void get_haplotype_peptides(const Haplotype& haplotype, set<string>& peptides);

class DiffProcessor : public HlaGroupsPairProcessor
{
public:
    DiffProcessor() : verbose(false) {}
private:
    string group1;
    string group2;
    vector<string> columns;
    vector<string> rows;

    bool verbose;

    struct CellData
    {
        size_t only1;
        size_t common;
        size_t only2;

        void set_data(size_t only1, size_t common, size_t only2)
        {
            this->only1  = only1;
            this->common = common;
            this->only2  = only2;
        }
    };

    struct HaplotypeData
    {
        map<string, map<string, CellData> > matrix;
    };

    map<string, HaplotypeData> haplotypes;

    void generate_headers(Sheet& sheet) const;
    void generate_sheet(Spreadsheet& spreadsheet, const string& haplotype, const HaplotypeData& haplotypeData) const;

    void set_verbose(bool value)
    {
        verbose = value;
    }

    void set_groups(const string& group1, const string& group2)
    {
        this->group1 = group1;
        this->group2 = group2;

        if (find(rows.begin(), rows.end(), group1) == rows.end())
            rows.push_back(group1);

        if (find(columns.begin(), columns.end(), group2) == columns.end())
            columns.push_back(group2);

        if (verbose)
            cerr << group1 << " / " << group2 << endl;
    }

    void add_haplotype(const string& haplotype, size_t only1, size_t common, size_t only2)
    {
        haplotypes[haplotype].matrix[group1][group2].set_data(only1, common, only2);

        if (verbose)
            cerr << "\t" << haplotype << ": [" << only1 << "," << common << "," << only2 << "]" << endl;
    }

    virtual ProcessorPriority get_priority() const  {   return 1.0f;    }
    virtual void generate_results(Spreadsheet& spreadsheet);
    virtual void process(const Group& g1, const Group& g2);

    void compare_haplotypes(const Haplotype& haplotype1, const Haplotype& haplotype2);
};

REGISTER_PROCESSOR(DiffProcessor);


template <class T>
static void check_all_in(const T& container2, const T& container1)
{
    // this is O(n^2), could be improved with a set.

    for (typename T::const_iterator it = container2.begin(); it != container2.end(); ++it)
    {
        if (find(container1.begin(), container1.end(), *it) == container1.end())
            cerr << "ERROR: Element " << it->name << " not found in counterpart!" << endl;
    }
}


template <class T>
static size_t sets_difference_count(const set<T>& set1, const set<T>& set2)
{
    set<T> only1;

    set_difference(set1.begin(), set1.end(), set2.begin(), set2.end(), inserter(only1, only1.begin()));

    return only1.size();
}

template <class T>
static size_t sets_common_count(const set<T>& set1, const set<T>& set2)
{
    set<T> common;

    set_intersection(set1.begin(), set1.end(), set2.begin(), set2.end(), inserter(common, common.begin()));

    return common.size();
}

void DiffProcessor::process(const Group& group1, const Group& group2)
{
    set_groups(group1.name, group2.name);

    // for each haplotype in group1, search it's equivalent in group2, and process:
    for (list<Haplotype>::const_iterator haplotype1 = group1.haplotypes.begin(); haplotype1 != group1.haplotypes.end(); ++haplotype1)
    {
        // search equivalent haplotype in group2:
        // TODO: REIMPLEMENT USING std::map (needs FXP support)
        const list<Haplotype>::const_iterator haplotype2 = find(group2.haplotypes.begin(), group2.haplotypes.end(), *haplotype1);

        if (haplotype2 != group2.haplotypes.end())
            compare_haplotypes(*haplotype1, *haplotype2);
        else
            cerr << "ERROR: Haplotype " << haplotype1->name << " not found in group " << group2.name << endl;
    }

    check_all_in(group2.haplotypes, group1.haplotypes);
}

void DiffProcessor::compare_haplotypes(const Haplotype& haplotype1, const Haplotype& haplotype2)
{
    set<string> peptides1;
    set<string> peptides2;

    // for each sequence in haplotype1, get the set of peptides. Same for haplotype2:
    get_haplotype_peptides(haplotype1, peptides1);
    get_haplotype_peptides(haplotype2, peptides2);

    add_haplotype(haplotype1.name,
        sets_difference_count(peptides1, peptides2),
        sets_common_count(peptides1, peptides2),
        sets_difference_count(peptides2, peptides1)
    );
}

static void get_haplotype_peptides(const Haplotype& haplotype, set<string>& peptides)
{
    for (list<Sequence>::const_iterator sequence = haplotype.sequences.begin(); sequence != haplotype.sequences.end(); ++sequence)
        peptides.insert(sequence->peptides.begin(), sequence->peptides.end());
}

void DiffProcessor::generate_headers(Sheet& sheet) const
{
    // first header row:
    {
        Row header1(sheet);
        const Style center( Style::ALIGN_CENTER );

        header1 << "" << v_separator;
        for (vector<string>::const_reverse_iterator group = columns.rbegin();
             group != columns.rend();
             ++group)
        {
            header1 << center << column_span(*group, 3) << v_separator;
        }
    }

    // second header row:
    {
        Row header2(sheet);
        const Style bottom_sep( Style::BORDER_BOTTOM | Style::ALIGN_CENTER );

        header2 << "" << v_separator;
        for (size_t group = 0; group < columns.size(); ++group)
            header2 << bottom_sep << "only in 1"
                    << bottom_sep << "common"
                    << bottom_sep << "only in 2"
                    << v_separator;
    }
}

void DiffProcessor::generate_sheet(Spreadsheet& spreadsheet, const string& haplotype, const HaplotypeData& haplotypeData) const
{
    Sheet new_sheet( spreadsheet, haplotype );
    const Style bottom_sep( Style::BORDER_TOP );

    generate_headers(new_sheet);

    for (vector<string>::const_iterator row = rows.begin();
         row != rows.end();
         ++row)
    {
        Row data(new_sheet);
        data << *row << v_separator;
        for (vector<string>::const_reverse_iterator column = columns.rbegin();
             column != columns.rend();
             ++column)
        {
            const map<string, map<string, CellData> >::const_iterator cell_row = haplotypeData.matrix.find(*row);

            if (cell_row != haplotypeData.matrix.end())
            {
                const map<string, CellData>::const_iterator cell = cell_row->second.find(*column);
                if (cell != cell_row->second.end())
                    data << cell->second.only1 << cell->second.common << cell->second.only2 << v_separator;
                else
                    data << "" << "" << "" << v_separator;
            }
            else
                cerr << "Internal error: " << *row << " not found in " << haplotype << endl;
        }
    }
}

void DiffProcessor::generate_results(Spreadsheet& spreadsheet)
{
    for (map<string, HaplotypeData>::const_iterator haplo = haplotypes.begin(); haplo != haplotypes.end(); ++haplo)
        generate_sheet(spreadsheet, haplo->first, haplo->second);
}

